/*
 * test_node.cpp
 *
 *  Created on: 16.09.2013
 *      Author: alex
 */



#include <bullet_rviz_test/stability_checker.h>
#include <stdlib.h>
#include <math.h>
#include <boost/make_shared.hpp>
#include <ros/ros.h>
#include <bullet/btBulletDynamicsCommon.h>
#include <bullet/btBulletCollisionCommon.h>
#include <bullet_rviz_test/rviz_motion_state.h>
#include <bullet_rviz_test/physics.h>
#include <bullet_rviz_test/physical_object.h>
#include <bullet_rviz_test/heightfield.h>
#include <bullet_rviz_test/heightfield_visualization.h>
#include <bullet_rviz_test/robot.h>
#include <bullet_rviz_test/robot_visualization.h>

static const uint32_t GROUND_COLOR = 0x999988FFU;
static const float WORST_DISTANCE = 0.2;
static const float WORST_ROTATION = 3.1415f / 4;
static const double SIMULATION_STEP_DURATION = 1.0/30.0;

using namespace octomap;
using namespace std;

namespace bullet_rviz_test
{

StabilityChecker::StabilityChecker(bool show_visualization_, bool load_from_file)
{
  max_tilt = 30.0 /M_PI*180.0; // [rad]
  max_movement = 0.15; // [m]

  last_accepted_octomap_time = ros::Time(0.0);

  show_visualization = show_visualization_;
  physics_ = boost::make_shared<Physics>();

  tree = NULL;
  if (load_from_file)
  {
    tree = new OcTree("data/ssrr_arena_map_3.bt");
    createHeightfieldFromOctomap();
  }
  else
  {
    sub = nh_.subscribe("/octomap_binary", 1, &StabilityChecker::map_callback, this);
  }

  RobotInfoPtr robot_info = boost::make_shared<RobotInfo>();

  robot_info->robot_ = boost::make_shared<Robot>();
  physics_->addPhysicalObject(robot_info->robot_);
//  robot_info->robot_->registerMotionCallback(
//      PhysicalObject::MotionCallback(
//          boost::bind(&StabilityChecker::robotMotionCallback, this,
//                      robot_info.get())));

  robot_info->visualization_ = boost::make_shared<RobotVisualization>(
      robot_info->robot_);

  robots_.push_back(robot_info);
}

bool StabilityChecker::check(double start_x, double start_y, double start_theta)
{
  ros::Rate loop_rate(30);

  ros::Time time_before_check = ros::Time::now();

  btTransform world_transform;

  bool had_impact = false;
  bool first_step = true;

//  double start_x, start_y;
  double last_x, last_y, last_z;

  double eps = 0.001;
  int num_iterations = 0;

  bool stable_configuration = false;

  resetRobots(start_x, start_y, start_theta);

  while (ros::ok() && !stable_configuration)
  {
	if (num_iterations > 100) {
//		std::cout << "computation took " << ((ros::Time::now() - time_before_check).toNSec()/1000000000.0) << "secs" << std::endl;
		return false;
	}

    if (show_visualization)
    {
      ros::spinOnce();
      loop_rate.sleep();
    }
    physics_->stepSimulation(SIMULATION_STEP_DURATION);

    robots_[0]->robot_->getRigidBody()->getMotionState()->getWorldTransform(world_transform);

    double x = world_transform.getOrigin().x();
    double y = world_transform.getOrigin().y();
    double z = world_transform.getOrigin().z();

    if (first_step)
    {
      first_step = false;

      last_x = world_transform.getOrigin().x();
      last_y = world_transform.getOrigin().y();
      last_z = world_transform.getOrigin().z();
    }
    else
    {
      if (fabs(last_x-x) < eps && fabs(last_y-y) < eps && fabs(last_z-z) < eps)
      {
        ROS_DEBUG("STATIC NOW!");
        stable_configuration = true;
      }
    }

    ROS_DEBUG("x: %f, y: %f, z: %f", x, y, z);

    if ((fabs(start_x - x) > eps || fabs(start_y - y) > eps) && !had_impact)
    {
      had_impact = true;
      ROS_DEBUG("IMPACT");

      // RESET SIMULATION...
      resetRobots(start_x, start_y, start_theta, z + 0.2);
    }

    last_x = x;
    last_y = y;
    last_z = z;

  }

  if (show_visualization)
  {
    heightfield_visualization_.reset();
  }

  // DETERMINE RETURN VALUE based on feasibilty
  world_transform = robots_[0]->robot_->getRigidBody()->getWorldTransform();
  btVector3 position = world_transform.getOrigin();

  btScalar moved_distance = position.distance(robots_[0]->initial_position_);

//  float cost = moved_distance / WORST_DISTANCE
//      + fabs(world_transform.getRotation().getAngle()) / WORST_ROTATION;

  // calculate angle between (0, 0, 1) and axis of angle-axis representation of rotation
  double angle = acos(world_transform.getRotation().getAxis().z());

  ROS_DEBUG("tilt [deg]: %f", angle/M_PI*180.0);
  ROS_DEBUG("moved distance [m]: %f", moved_distance);

//  std::cout << "computation took " << ((ros::Time::now() - time_before_check).toNSec()/1000000000.0) << "secs" << std::endl;
  return (fabs(angle) < max_tilt && moved_distance < max_movement);
}

void StabilityChecker::map_callback(const octomap_msgs::OctomapBinary::ConstPtr& msg)
{
  if (ros::Time::now() - last_accepted_octomap_time < ros::Duration(20))
	  return;

  last_accepted_octomap_time = ros::Time::now();

  ros::Time time_before_map = ros::Time::now();

  tree = octomap_msgs::binaryMsgDataToMap(msg->data);

  createHeightfieldFromOctomap();

//  std::cout << "loading map took " << ((ros::Time::now() - time_before_map).toNSec()/1000000000.0) << "secs" << std::endl;

  // we no longer need the tree
  delete tree;
}

void StabilityChecker::resetRobots(double robot_x, double robot_y, double robot_theta, double drop_z)
{
  // set robot position

  btVector3 heightfield_extents = heightfield_->getExtents();
  btVector3 heightfield_position =
      heightfield_->getRigidBody()->getWorldTransform().getOrigin();

  RobotPtr & robot = robots_[0]->robot_;

  double start_z;
  if (drop_z < -500.0)
    start_z = heightfield_extents.z() + 1.0;
  else
    start_z = drop_z;


  btScalar x = heightfield_position.x() + robot_x;
  btScalar y = heightfield_position.y() + robot_y;
  btScalar z = heightfield_position.z() + start_z;

  btVector3 position(x, y, z);
  robot->getRigidBody()->setWorldTransform(
      btTransform(btQuaternion(0.0, 0.0, robot_theta), position));
  robot->getRigidBody()->setLinearVelocity(btVector3(0, 0, 0));
  robot->getRigidBody()->setAngularVelocity(btVector3(0, 0, 0));
  robot->getRigidBody()->activate();

  robots_[0]->initial_position_ = btVector3(x, y,
                                            z - robot->getHalfExtents().z());

}

bool StabilityChecker::has_map()
{
  return (tree != NULL);
}



void StabilityChecker::createHeightfieldFromOctomap()
{
  double min_x, min_y, min_z;
  tree->getMetricMin(min_x, min_y, min_z);

  double max_x, max_y, max_z;
  tree->getMetricMax(max_x, max_y, max_z);

  btVector3 heightfield_extents(max_x-min_x, max_y-min_y, max_z-min_z);

  double node_size = tree->getResolution();

  int width = ceil((max_x - min_x)/node_size);
  int height = ceil((max_y - min_y)/node_size);

  // create height map
  float * data = new float[width * height];
  for (int i=0; i<width*height; i++)
    data[i] = min_z;

  for(OcTree::leaf_iterator it = tree->begin_leafs(),
         end=tree->end_leafs(); it!= end; ++it)
  {
    point3d center = it.getCoordinate();
    double current_size = it.getSize();

    double left_upper_x = center.x() - current_size/2.0 + node_size/2.0;
    double left_upper_y = center.y() - current_size/2.0 + node_size/2.0;

    for (double x = left_upper_x; x < left_upper_x + current_size; x += node_size)
    {
      for (double y = left_upper_y; y < left_upper_y + current_size; y += node_size)
      {
        int index_x = floor((x-min_x)/node_size);
        int index_y = floor((y-min_y)/node_size);

        double block_top = center.z() + node_size/2.0;

        // value > 0 means more than 50% (prob=exp(value)/(1+exp(value)))
        if (data[index_x + index_y*width] < block_top && it->getValue() > 0.0)
          data[index_x + index_y*width] = block_top;
      }
    }
  }

  if (heightfield_ != NULL)
  {
    physics_->removePhysicalObject(heightfield_);

    std::cout << "remove object" << std::endl;

    heightfield_.reset();
  }

  //heightfield_ = boost::make_shared<Heightfield>(heightfield_extents, data,
  //                                               width, height);

  heightfield_.reset(new Heightfield(heightfield_extents, data, width, height));

  physics_->addPhysicalObject(heightfield_);

  heightfield_visualization_ = boost::make_shared<HeightfieldVisualization>(
      heightfield_, GROUND_COLOR);

  // display map and wait for node to be connected
  if (show_visualization)
  {
    ros::Duration(0.5).sleep();
    heightfield_visualization_->visualize();
  }


  delete[] data;
}

}
